#!/bin/bash
#
# Disk Retry QA Script -- Real Datadog Intake via Local Proxy
#
# Runs a local TCP proxy (localhost:18443) that forwards to Datadog's intake.
# The agent is configured to send through this proxy. Stopping/starting the
# proxy simulates a clean network outage.
#
# Flow:
#   1. Start proxy, agent sends logs through it     -> verify in Log Explorer
#   2. Kill proxy (simulates outage)                 -> payloads go to disk
#   3. Stop log writer, let pipeline drain
#   4. Restart proxy                                 -> payloads replay to Datadog
#
# Prerequisites:
#   - Build the agent: dda inv agent.build --build-exclude=systemd
#   - Valid API key in dev/dist/datadog.yaml
#   - datadog.yaml must set: logs_config.logs_dd_url: "localhost:18443"
#   - disk_retry config enabled
#
# Usage:
#   chmod +x dev/disk-retry-qa-real.sh
#   ./dev/disk-retry-qa-real.sh

AGENT_BIN="./bin/agent/agent"
LOG_FILE="/tmp/test-disk-retry-real.log"
RETRY_DIR="/tmp/dd-logs-retry"
PROXY_PORT=18443
INTAKE_HOST="agent-http-intake.logs.datadoghq.com"
INTAKE_PORT=443
LOG_WRITER_PID=""
PROXY_PID=""

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

cleanup() {
    echo -e "\n${YELLOW}Cleaning up...${NC}"
    [ -n "$LOG_WRITER_PID" ] && kill "$LOG_WRITER_PID" 2>/dev/null
    [ -n "$PROXY_PID" ] && kill "$PROXY_PID" 2>/dev/null
    sleep 1
    echo "Done."
}
trap cleanup EXIT

# ── Helpers ──────────────────────────────────────────────────────────

start_proxy() {
    # TCP proxy: accepts TLS connections on localhost:18443, forwards to Datadog intake.
    # Uses socat if available, falls back to python.
    if command -v socat &>/dev/null; then
        socat TCP-LISTEN:$PROXY_PORT,fork,reuseaddr OPENSSL:$INTAKE_HOST:$INTAKE_PORT,verify=0 &
        PROXY_PID=$!
    else
        python3 -c "
import socket, ssl, threading, sys

def forward(src, dst):
    try:
        while True:
            data = src.recv(65536)
            if not data:
                break
            dst.sendall(data)
    except:
        pass
    finally:
        src.close()
        dst.close()

def handle(client):
    ctx = ssl.create_default_context()
    upstream = ctx.wrap_socket(socket.socket(), server_hostname='$INTAKE_HOST')
    upstream.connect(('$INTAKE_HOST', $INTAKE_PORT))
    t1 = threading.Thread(target=forward, args=(client, upstream), daemon=True)
    t2 = threading.Thread(target=forward, args=(upstream, client), daemon=True)
    t1.start()
    t2.start()
    t1.join()
    t2.join()

server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(('127.0.0.1', $PROXY_PORT))
server.listen(32)
print(f'TCP proxy listening on :$PROXY_PORT -> $INTAKE_HOST:$INTAKE_PORT', flush=True)
while True:
    client, addr = server.accept()
    threading.Thread(target=handle, args=(client,), daemon=True).start()
" &
        PROXY_PID=$!
    fi
    sleep 2
    echo -e "${GREEN}Proxy started on :$PROXY_PORT -> $INTAKE_HOST:$INTAKE_PORT${NC}"
}

stop_proxy() {
    if [ -n "$PROXY_PID" ]; then
        kill "$PROXY_PID" 2>/dev/null
        sleep 1
        PROXY_PID=""
    fi
}

start_log_writer() {
    local tag=$1
    (
        i=0
        while true; do
            echo "{\"message\": \"disk-retry-qa $tag line $i\", \"timestamp\": \"$(date -u +%Y-%m-%dT%H:%M:%SZ)\", \"phase\": \"$tag\"}" >> "$LOG_FILE"
            i=$((i + 1))
            sleep 1
        done
    ) &
    LOG_WRITER_PID=$!
}

stop_log_writer() {
    if [ -n "$LOG_WRITER_PID" ]; then
        kill "$LOG_WRITER_PID" 2>/dev/null
        sleep 1
        LOG_WRITER_PID=""
    fi
}

retry_file_count() {
    if [ -d "$RETRY_DIR" ]; then
        sudo find "$RETRY_DIR" -name "*.retry" 2>/dev/null | wc -l | tr -d ' '
    else
        echo "0"
    fi
}

retry_dir_size() {
    if [ -d "$RETRY_DIR" ]; then
        sudo du -sh "$RETRY_DIR" 2>/dev/null | cut -f1
    else
        echo "0"
    fi
}

# ── Main ─────────────────────────────────────────────────────────────

echo "============================================"
echo "  Disk Retry QA -- Real Datadog via Proxy"
echo "============================================"
echo ""

# Clean state
rm -f "$LOG_FILE" 2>/dev/null
sudo rm -rf "$RETRY_DIR" 2>/dev/null
touch "$LOG_FILE"

# Check agent binary
if [ ! -x "$AGENT_BIN" ]; then
    echo -e "${RED}Agent binary not found at $AGENT_BIN${NC}"
    echo "Build it first: dda inv agent.build --build-exclude=systemd"
    exit 1
fi

# ── Phase 1: Normal operation ────────────────────────────────────────
echo -e "${GREEN}=== Phase 1: Normal operation ===${NC}"
echo "Starting proxy..."
start_proxy

echo "Starting log writer (1 line/sec, tagged phase=pre-outage)..."
start_log_writer "pre-outage"

echo ""
echo -e "${YELLOW}Start the agent in another terminal:${NC}"
echo ""
echo "  cd $(pwd)"
echo "  sudo $AGENT_BIN run -c dev/dist/datadog.yaml"
echo ""
echo "Check Datadog Log Explorer for: service:disk-retry-test"
echo ""
read -p "Press ENTER once the agent is running and you see logs in Datadog... "

echo ""
echo -e "${CYAN}Retry dir: $(retry_file_count) files, $(retry_dir_size)${NC}"

# ── Phase 2: Simulate outage (45 seconds) ───────────────────────────
echo ""
echo -e "${RED}=== Phase 2: Killing proxy (45s outage) ===${NC}"
stop_log_writer
start_log_writer "during-outage"

stop_proxy
echo "Proxy killed. Agent cannot reach Datadog."
echo "Logs tagged 'during-outage' being written. Payloads should go to disk."
echo ""

for i in $(seq 1 45); do
    count=$(retry_file_count)
    size=$(retry_dir_size)
    printf "  [%2ds] retry files: %-6s disk usage: %s\n" "$i" "$count" "$size"
    sleep 1
done

PHASE2_FINAL=$(retry_file_count)
PHASE2_SIZE=$(retry_dir_size)
echo ""
if [ "$PHASE2_FINAL" -gt 0 ] 2>/dev/null; then
    echo -e "${GREEN}Phase 2 result: $PHASE2_FINAL files, $PHASE2_SIZE on disk${NC}"
else
    echo -e "${RED}Phase 2 result: NO retry files -- check agent logs${NC}"
fi

# ── Phase 3: Stop log writer ─────────────────────────────────────────
echo ""
echo -e "${YELLOW}=== Phase 3: Stopping log writer ===${NC}"
stop_log_writer
echo "Waiting 15s for pipeline to flush to disk..."
echo ""

for i in $(seq 1 15); do
    count=$(retry_file_count)
    size=$(retry_dir_size)
    printf "  [%2ds] retry files: %-6s disk usage: %s\n" "$i" "$count" "$size"
    sleep 1
done

PRE_RECOVERY=$(retry_file_count)
PRE_RECOVERY_SIZE=$(retry_dir_size)
echo ""
echo -e "${CYAN}Files before recovery: $PRE_RECOVERY ($PRE_RECOVERY_SIZE)${NC}"

# ── Phase 4: Recovery ────────────────────────────────────────────────
echo ""
echo -e "${GREEN}=== Phase 4: Restarting proxy + new logs ===${NC}"
start_log_writer "post-recovery"
start_proxy

echo "Watching for drain (5 minutes max)..."
echo ""

LAST_COUNT=$PRE_RECOVERY
for i in $(seq 1 300); do
    count=$(retry_file_count)
    size=$(retry_dir_size)
    printf "  [%3ds] retry files: %-6s disk usage: %s\n" "$i" "$count" "$size"

    if [ "$count" = "0" ]; then
        echo ""
        echo -e "  ${GREEN}All retry files replayed at ${i}s!${NC}"
        break
    fi

    if [ "$count" -lt "$LAST_COUNT" ] 2>/dev/null; then
        echo -e "  ${CYAN}  ^ files decreasing -- replay active${NC}"
    fi
    LAST_COUNT=$count

    sleep 1
done

# ── Summary ──────────────────────────────────────────────────────────
echo ""
echo "============================================"
echo "  QA Summary"
echo "============================================"
FINAL_COUNT=$(retry_file_count)
FINAL_SIZE=$(retry_dir_size)
echo "  Phase 2 (outage):   $PHASE2_FINAL files ($PHASE2_SIZE)"
echo "  Pre-recovery:       $PRE_RECOVERY files ($PRE_RECOVERY_SIZE)"
echo "  Final:              $FINAL_COUNT files ($FINAL_SIZE)"
echo ""
if [ "$PHASE2_FINAL" -gt 0 ] 2>/dev/null && [ "$FINAL_COUNT" -lt "$PRE_RECOVERY" ] 2>/dev/null; then
    echo -e "  ${GREEN}PASS: Payloads stored and replaying to Datadog${NC}"
elif [ "$FINAL_COUNT" = "0" ]; then
    echo -e "  ${GREEN}PASS: All payloads replayed to Datadog${NC}"
else
    echo -e "  ${YELLOW}PARTIAL: Replay in progress, may need more time${NC}"
fi
echo ""
echo "Check Datadog Log Explorer for all three phases:"
echo "  service:disk-retry-test @phase:pre-outage"
echo "  service:disk-retry-test @phase:during-outage  (replayed from disk)"
echo "  service:disk-retry-test @phase:post-recovery"
echo ""
echo "You can now stop the agent."